// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/aggregate.h>
#include <pollux/exec/simple_aggregate_adapter.h>
#include <pollux/expression/function_signature.h>
#include <pollux/expression/vector_writers.h>

using namespace kumo::pollux::exec;

namespace kumo::pollux::aggregate {

namespace {

// Implementation of the average aggregation function through the
// SimpleAggregateAdapter.
template <typename T>
class AverageAggregate {
 public:
  // Type(s) of input vector(s) wrapped in Row.
  using InputType = Row<T>;

  // Type of intermediate result vector wrapped in Row.
  using IntermediateType =
      Row</*sum*/ double,
          /*count*/ int64_t>;

  // Type of output vector.
  using OutputType =
      std::conditional_t<std::is_same_v<T, float>, float, double>;

  static bool toIntermediate(
      exec::out_type<Row<double, int64_t>>& out,
      exec::arg_type<T> in) {
    out.copy_from(std::make_tuple(static_cast<double>(in), 1));
    return true;
  }

  struct AccumulatorType {
    double sum_;
    int64_t count_;

    AccumulatorType() = delete;

    // Constructor used in initializeNewGroups().
    explicit AccumulatorType(
        HashStringAllocator* /*allocator*/,
        AverageAggregate* /*fn*/) {
      sum_ = 0;
      count_ = 0;
    }

    // addInput expects one parameter of exec::arg_type<T> for each child-type T
    // wrapped in InputType.
    void addInput(HashStringAllocator* /*allocator*/, exec::arg_type<T> data) {
      sum_ += data;
      count_ = checkedPlus<int64_t>(count_, 1);
    }

    // combine expects one parameter of exec::arg_type<IntermediateType>.
    void combine(
        HashStringAllocator* /*allocator*/,
        exec::arg_type<Row<double, int64_t>> other) {
      // Both field of an intermediate result should be non-null because
      // writeIntermediateResult() never make an intermediate result with a
      // single null.
      POLLUX_CHECK(other.at<0>().has_value());
      POLLUX_CHECK(other.at<1>().has_value());
      sum_ += other.at<0>().value();
      count_ = checkedPlus<int64_t>(count_, other.at<1>().value());
    }

    bool writeFinalResult(exec::out_type<OutputType>& out) {
      out = sum_ / count_;
      return true;
    }

    bool writeIntermediateResult(exec::out_type<IntermediateType>& out) {
      out = std::make_tuple(sum_, count_);
      return true;
    }
  };
};

} // namespace

exec::AggregateRegistrationResult registerSimpleAverageAggregate(
    const std::string& name) {
  std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures;

  for (const auto& inputType : {"smallint", "integer", "bigint", "double"}) {
    signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                             .returnType("double")
                             .intermediateType("row(double,bigint)")
                             .argumentType(inputType)
                             .build());
  }

  signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                           .returnType("real")
                           .intermediateType("row(double,bigint)")
                           .argumentType("real")
                           .build());

  return exec::registerAggregateFunction(
      name,
      std::move(signatures),
      [name](
          core::AggregationNode::Step step,
          const std::vector<TypePtr>& argTypes,
          const TypePtr& resultType,
          const core::QueryConfig& /*config*/)
          -> std::unique_ptr<exec::Aggregate> {
        POLLUX_CHECK_LE(
            argTypes.size(), 1, "{} takes at most one argument", name);
        auto inputType = argTypes[0];
        if (exec::isRawInput(step)) {
          switch (inputType->kind()) {
            case TypeKind::SMALLINT:
              return std::make_unique<
                  SimpleAggregateAdapter<AverageAggregate<int16_t>>>(
                  step, argTypes, resultType);
            case TypeKind::INTEGER:
              return std::make_unique<
                  SimpleAggregateAdapter<AverageAggregate<int32_t>>>(
                  step, argTypes, resultType);
            case TypeKind::BIGINT:
              return std::make_unique<
                  SimpleAggregateAdapter<AverageAggregate<int64_t>>>(
                  step, argTypes, resultType);
            case TypeKind::REAL:
              return std::make_unique<
                  SimpleAggregateAdapter<AverageAggregate<float>>>(
                  step, argTypes, resultType);
            case TypeKind::DOUBLE:
              return std::make_unique<
                  SimpleAggregateAdapter<AverageAggregate<double>>>(
                  step, argTypes, resultType);
            default:
              POLLUX_FAIL(
                  "Unknown input type for {} aggregation {}",
                  name,
                  inputType->kindName());
          }
        } else {
          switch (resultType->kind()) {
            case TypeKind::REAL:
              return std::make_unique<
                  SimpleAggregateAdapter<AverageAggregate<float>>>(
                  step, argTypes, resultType);
            case TypeKind::DOUBLE:
            case TypeKind::ROW:
              return std::make_unique<
                  SimpleAggregateAdapter<AverageAggregate<double>>>(
                  step, argTypes, resultType);
            default:
              POLLUX_FAIL(
                  "Unsupported result type for final aggregation: {}",
                  resultType->kindName());
          }
        }
      },
      true /*registerCompanionFunctions*/,
      true /*overwrite*/);
}

} // namespace kumo::pollux::aggregate
